pytorch模型预测 |
您所在的位置：网站首页 › pytorch train函数 › pytorch模型预测 |
pytorch模型预测
|
pytorch模型预测_函数整理——predict.py
一、预测_pred1.1 模型全部加载——torch.load()1.2 模型参数加载——model.load_state_dict()1.3 设置训练/评估模式——model.train()和model.eval()1.4 数组转换成张量——torch.from_numpy()1.5 数据从其他设备转移到cpu上—— tensor.cpu()1.6 取tensor中的数据—— tensor.data1.7 np.array(pred.data.cpu()[0])[0]
一、预测_pred
1、选择设备，加载模型，设置模型测试（部分函数见：pytorch模型训练函数整理——train.py） 2、读取图片，图片预处理（灰度图、尺寸、通道)，转为张量tensor并拷贝到设备中 3、放入模型预测 4、提取预测结果，对结果进行后处理（从概率处理成像素），保存图片 1.1 模型全部加载——torch.load() 作用：用来加载整个模型（结构+参数）；torch.save()函数一般在train.py文件中可以找到。主要参数： torch.load(f,map_location=) 参数含义f加载模型的路径map_location模型存储位置；如device 1.2 模型参数加载——model.load_state_dict() 作用：加载不带结构的模型参数应用： state_dict = torch.load('model_state_dict.pth') model.load_state_dict(state_dict) 1.3 设置训练/评估模式——model.train()和model.eval() 用法：如果模型中有BN层（Batch Normalization）和 Dropout层 ，需要在训练开始之前写上 model.trian() ，在测试时写上 model.eval()作用： 1）model.train() 是保证 BN 层用每一批数据的均值和方差 2） model.eval() 是保证 BN 用全部训练数据的均值和方差； 1.4 数组转换成张量——torch.from_numpy()作用：把数组转换成张量，且二者共享内存 1.5 数据从其他设备转移到cpu上—— tensor.cpu()将数据的处理设备从其他设备（例如cuda)拿到cpu上 1.6 取tensor中的数据—— tensor.data只取出原来Tensor的tensor数据， 丢弃了grad、grad_fn等额外的信息 1.7 np.array(pred.data.cpu()[0])[0]pred是模型对图像的分割预测数据 操作目的是取出二维数组形式的数据 pred # 原数据 pred.data # 丢弃grad_fn pred.data.cpu() # 从cuda转入cpu pred.data.cpu()[0] # 取第一个内容，即数据内容 np.array(pred.data.cpu()[0])[0] # 用数组存储tensor中的第一维，去除张量中多余的一个[] |
【本文地址】
今日新闻 |
推荐新闻 |